#!/usr/bin/env python3

# Retrieval Grader

import torch

from pydantic import BaseModel
from typing import Annotated
from typing_extensions import TypedDict

from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langchain_community.chat_models.outlines import ChatOutlines
from langchain_core.messages import HumanMessage

import json

from langchain_core.messages import ToolMessage


class BasicToolNode:
    """A node that runs the tools requested in the last AIMessage."""

    def __init__(self, tools: list) -> None:
        self.tools_by_name = {tool.name: tool for tool in tools}

    def __call__(self, inputs: dict):
        if messages := inputs.get("messages", []):
            message = messages[-1]
        else:
            raise ValueError("No message found in input")
        outputs = []
        for tool_call in message.tool_calls:
            tool_result = self.tools_by_name[tool_call["name"]].invoke(
                tool_call["args"]
            )
            outputs.append(
                ToolMessage(
                    content=json.dumps(tool_result),
                    name=tool_call["name"],
                    tool_call_id=tool_call["id"],
                )
            )
        return {"messages": outputs}

def get_model():
    """define for case of caching"""
    # model = models.transformers(
    #     "gpt2",               # Model name
    #     max_length=100,       # Generate up to 100 tokens
    #     temperature=0.7,      # Moderate randomness
    #     top_p=0.9,            # Nucleus sampling
    #     repetition_penalty=1.2,  # Penalize repetition
    #     stop=["\n"],          # Stop on newline character
    #     device="cuda"         # Use GPU
    # )
    model_path = "/home/wangxianda/model_zoo/llama_3.2_3b_instruct"
    model = ChatOutlines(model=model_path)
    model.client.model.to("cuda", dtype=torch.float16)
    return model

_set_env("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")
from langchain_community.tools.tavily_search import TavilySearchResults

search = TavilySearchResults(max_results=2)
tools = [search]
tool_node = BasicToolNode(tools=tools)
model = get_model()
model_with_tools = model.bind_tools(tools)

class State(TypedDict):
    # Messages have the type "list". The `add_messages` function
    # in the annotation defines how this state key should be updated
    # (in this case, it appends messages to the list, rather than overwriting them)
    messages: Annotated[list, add_messages]


graph_builder = StateGraph(State)

def chatbot(state: State):
    return {"messages": [model.invoke(state["messages"])]}


graph_builder.add_node("tools", tool_node)
graph_builder.add_node("chatbot", chatbot)
graph_builder.add_edge(START, "chatbot")
graph_builder.add_edge("chatbot", END)
graph = graph_builder.compile()

graph.get_graph().print_ascii()


def stream_graph_updates(user_input: str):
    for event in graph.stream({"messages": [("user", user_input)]}):
        for value in event.values():
            print("Assistant:", value["messages"][-1].content)


while True:
    try:
        user_input = input("User: ")
        if user_input.lower() in ["quit", "exit", "q"]:
            print("Goodbye!")
            break

        stream_graph_updates(user_input)
    except:
        # fallback if input() is not available
        user_input = "What do you know about LangGraph?"
        print("User: " + user_input)
        stream_graph_updates(user_input)
        break
# messages = [HumanMessage(content="What will the capital of mars be called?")]
# response = model.invoke(messages)

# print( response.content )

# class Tool(BaseModel):
#     name: str
#     tool_input: str
#     cycles: int

# model.json_schema = Tool
# response = model.invoke("Who are the main contributors to LangChain?")
# person = Tool.model_validate_json(response.content)

# print(person)
